import base64
import cv2
import imutils
import numpy as np


class PageCorrection:
    def __init__(self):
        self.__processingPage = None
        self.__debugPic = None
        self.pages = None
        self.h = 0
        self.w = 0

    def getDebugPic(self):
        return self.__debugPic.copy()

    def setDebugPic(self, img):
        self.__debugPic = img

    def getProcessingPage(self):
        return self.__processingPage.copy()

    def setProcessingPage(self, img):
        self.__processingPage = img
        self.h, self.w = img.shape[0:2]

    def imshow(self, img):
        cv2.imshow("imShow", imutils.resize(img, height=700))
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def getThres(self, img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        blur = cv2.GaussianBlur(gray, (3, 3), 0)
        _, thres = cv2.threshold(blur, 127, 255, cv2.THRESH_BINARY)
        return thres

    def curvePage(self, img, contour):
        mask = np.zeros_like(img)
        cv2.drawContours(mask, [contour], -1, (255, 255, 255), cv2.FILLED)
        res = cv2.bitwise_and(img, mask)
        x, y, w, h = cv2.boundingRect(cv2.cvtColor(res, cv2.COLOR_RGB2GRAY))
        res = res[y: y + h, x: x + w]
        return res

    def getPages(self, img, thres):
        ret = []
        cnts, _ = cv2.findContours(
            thres,
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE
        )
        cnts = sorted(cnts, key=cv2.contourArea, reverse=True)[:2]

        for (index, item) in enumerate(cnts):
            if cv2.contourArea(item) < (
                    img.shape[0] * img.shape[1]
            ) * 0.25:
                cnts.pop(index)
            else:
                ret.append(
                    (self.curvePage(img, item), sum([i[0][0] for i in item])))

        if len(ret) == 2 and ret[0][1] > ret[1][1]:
            return [ret[1][0], ret[0][0]]

        return [i[0] for i in ret]

    def getTextContourMask(self, img):
        _, w = img.shape[0: 2]
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        mask = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C,
            cv2.THRESH_BINARY_INV, 45, 25
        )
        mask = cv2.dilate(mask, np.ones((1, 9), dtype=np.uint8),
                          iterations=2 * int(w / 500))
        # mask = cv2.erode(mask, np.ones((3, 1), dtype=np.uint8))
        return mask

    # 求取版心外接四边形

    def getTypeArea(self, side):
        pp = self.getProcessingPage()
        contourRect = []
        xList = []
        yList = []
        mask = self.getTextContourMask(self.getProcessingPage())
        contours, _ = cv2.findContours(
            mask,
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_NONE
        )
        for item in contours:  # 筛选contours
            xmin, ymin, width, height = cv2.boundingRect(item)
            if (width < 0.03 * self.w or height < 0.005 * self.h or
                    width < 1.5 * height):
                continue  # 抛弃太短或者太窄的边缘
            if ymin < 0.025 * self.h or ymin + height > self.h - 0.025 * self.h:
                continue  # 抛弃距离页边距太近的边缘
            if side == "left" and xmin < 0.02 * self.w:
                continue
            if side == "right" and xmin + width > self.w - 0.02 * self.w:
                continue
            contourRect.append((xmin, ymin, width, height))
            xList.extend((xmin, xmin + width))
            yList.extend((ymin, ymin + height))

            cv2.rectangle(
                pp, (xmin, ymin),
                (xmin + width, ymin + height),
                (0, 255, 0), 1
            )
        xList = sorted(xList)
        yList = sorted(yList)
        lt = (xList[0], yList[0])
        rt = (xList[-1], yList[0])
        lb = (xList[0], yList[-1])
        rb = (xList[-1], yList[-1])

        cv2.rectangle(pp, lt, rb, (0, 0, 255), 3)
        self.setDebugPic(pp)
        return lt, rt, lb, rb

    def getPageExtPoint(self, cnt):
        pp = self.getProcessingPage()
        extTop = tuple(cnt[cnt[:, :, 1].argmin()][0])
        extBot = tuple(cnt[cnt[:, :, 1].argmax()][0])

        cv2.circle(pp, extTop, 8, (0, 255, 0), 8)
        cv2.circle(pp, extBot, 8, (0, 255, 0), 8)
        self.setDebugPic(pp)

        return extTop, extBot

    def getDistortParameters(self, cnt, typeArea):
        pp = self.getProcessingPage()
        topList = []
        botList = []
        for item in cnt:
            if item[0][1] > 0.1 * self.h and item[0][1] < 0.9 * self.h:
                continue
            if item[0][0] < typeArea[0][0] or item[0][0] > typeArea[1][0]:
                continue
            if item[0][1] <= typeArea[0][1]:
                topList.append(item)
                cv2.circle(pp, item[0], 5, [0, 0, 255], -1)
            elif item[0][1] >= typeArea[2][1]:
                botList.append(item)
                cv2.circle(pp, item[0], 5, [0, 0, 255], -1)
        topCntMin = sorted(
            topList, key=lambda item: item[0][1] - typeArea[0][1], reverse=True)[0]
        botCntMin = sorted(
            botList, key=lambda item: item[0][1] - typeArea[2][1])[0]
        cv2.line(
            pp, (0, topCntMin[0][1]),
            (self.w, topCntMin[0][1]),
            [0, 255, 0], 2
        )
        cv2.line(
            pp, (0, botCntMin[0][1]),
            (self.w, botCntMin[0][1]),
            [0, 255, 0], 2
        )
        self.setDebugPic(pp)
        return topCntMin[0][1], botCntMin[0][1]

    def getPageContour(self):
        cnts, _ = cv2.findContours(
            self.getThres(self.getProcessingPage()),
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE
        )
        return sorted(cnts, key=cv2.contourArea, reverse=True)[0]

    def correction(self, extBot, extTop, means, typeArea):
        peekX = round((extBot[0] + extTop[0]) / 2)
        left = self.getProcessingPage()[:, 0: peekX]
        M = cv2.getPerspectiveTransform(
            np.float32([
                [0, 0],
                [0, self.h],
                [peekX, extTop[1]],
                [peekX, extBot[1]]
            ]),
            np.float32([
                [0, 0],
                [0, self.h],
                [peekX, means[0]],
                [peekX, means[1]]
            ])
        )
        left = cv2.warpPerspective(left, M, (peekX, self.h))
        right = self.getProcessingPage()[:, peekX:]
        M = cv2.getPerspectiveTransform(
            np.float32([
                [0, extTop[1]],
                [0, extBot[1]],
                typeArea[1],
                typeArea[3]
            ]),
            np.float32([
                [0, means[0]],
                [0, means[1]],
                typeArea[1],
                typeArea[3]
            ])
        )
        right = cv2.warpPerspective(right, M, (self.w - peekX, self.h))
        res = cv2.hconcat([left, right])
        M = cv2.getPerspectiveTransform(
            np.float32(typeArea),
            np.float32([[0, 0], [self.w, 0], [0, self.h], [self.w, self.h]])
        )
        res = cv2.warpPerspective(res, M, (self.w, self.h))
        self.setDebugPic(res)
        return res

    def enhance(self, img):
        res = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        close = cv2.morphologyEx(
            res, cv2.MORPH_CLOSE,
            cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
        )
        res = ~(close - res)
        self.setDebugPic(res)
        return res

    def openFromBase64(self, imgBase64):
        imgBase64 = str(imgBase64)
        b = imgBase64[imgBase64.find(",") + 1:]
        data = base64.b64decode(b)
        imgArray = np.fromstring(data, np.uint8)
        img = cv2.imdecode(imgArray, cv2.IMREAD_COLOR)
        thres = self.getThres(img)
        self.pages = self.getPages(img, thres)
        return self.pages

    def openFromFile(self, filename):
        img = cv2.imread(filename)
        thres = self.getThres(img)
        self.pages = self.getPages(img, thres)
        return self.pages

    def process(self):
        ret = []
        for page in self.pages:
            self.setProcessingPage(page)
            contour = self.getPageContour()
            extTop, extBot = self.getPageExtPoint(contour)
            side = "right" if extTop[0] + extBot[0] < self.w else "left"
            typeArea = self.getTypeArea(side)
            means = self.getDistortParameters(contour, typeArea)
            res = self.correction(extBot, extTop, means, typeArea)
            ret.append(self.enhance(res))
        return ret

    def debug(self):
        for page in self.pages:
            self.setProcessingPage(page)
            self.imshow(page)
            contour = self.getPageContour()
            extTop, extBot = self.getPageExtPoint(contour)
            side = "right" if extTop[0] + extBot[0] < self.w else "left"
            self.imshow(self.getDebugPic())
            typeArea = self.getTypeArea(side)
            self.imshow(self.getDebugPic())
            means = self.getDistortParameters(contour, typeArea)
            self.imshow(self.getDebugPic())
            res = self.correction(extBot, extTop, means, typeArea)
            self.imshow(self.getDebugPic())
            self.enhance(res)
            self.imshow(self.getDebugPic())


if __name__ == '__main__':
    o = PageCorrection()
    o.openFromFile("IMG_4719.JPG")
    o.debug()
